using System.Collections.Generic;
using System.Windows.Forms;

namespace LeiRay
{
    public class TreeViewState
    {
        private readonly TreeView treeView;
        Dictionary<string, bool> list = new Dictionary<string, bool>();
        private string selectedNode = null;

        public TreeViewState(TreeView treeView)
        {
            this.treeView = treeView;
        }

        public void SaveTreeView()
        {
            list.Clear();
            SaveTreeViewExpandedState(treeView.Nodes);
            if (treeView.SelectedNode != null)
                selectedNode = treeView.SelectedNode.FullPath;
        }

        public void RestoreTreeView()
        {
            RestoreTreeViewExpandedState(treeView.Nodes);
            if(selectedNode != null)
            {
                TreeNode findTreeNodeFromPath = FindTreeNodeFromPath(selectedNode, treeView.TopNode);
                if(findTreeNodeFromPath != null)
                {
                    treeView.SelectedNode = findTreeNodeFromPath;
                }
            }
        }

        private TreeNode FindTreeNodeFromPath(string path, TreeNode treeNode)
        {
            // Found it
            if(treeNode.FullPath == path)
            {
                return treeNode;
            }

            foreach (TreeNode node in treeNode.Nodes)
            {
                TreeNode findTreeNodeFromPath = FindTreeNodeFromPath(path, node);
                if(findTreeNodeFromPath != null)
                {
                    return findTreeNodeFromPath;
                }
            }

            return null;
        }

        private void SaveTreeViewExpandedState(TreeNodeCollection nodes)
        {
            foreach (TreeNode node in nodes)
            {
                list.Add(node.FullPath, node.IsExpanded);
                if (node.Nodes.Count > 0)
                {
                    SaveTreeViewExpandedState(node.Nodes);
                }
            }
        }

        private void RestoreTreeViewExpandedState(TreeNodeCollection nodes)
        {
            foreach (TreeNode node in nodes)
            {
                if(!list.ContainsKey(node.FullPath))
                {
                    continue;
                }

                if (list[node.FullPath])
                {
                    node.Expand();
                }
                if (node.Nodes.Count > 0)
                {
                    RestoreTreeViewExpandedState(node.Nodes);
                }
            }
        }
    }
}